import os
import struct
import time

import cv2
import numpy as np

time_data = None
# ***不修改的内容***
cam_x = 320
cam_y = 215
cam_h = 1.500000
ppd = 35.6
fx = 919.000000
fy = 920.000000
width = 1280
height = 720
frame_rate = 30
img_format = 'NV12'
img_size = 1382400
Odr_GYR = 400
Range_GYR = 125
Odr_ACC = 250
Range_ACC = 2

# ***修改的内容***
image_frame_idx = 0
image_exists = 0
imu_valid = 1
gyr_raw_x = 0
gyr_raw_y = 0
gyr_raw_z = 0
gyr_lf_x = 0
gyr_lf_y = 0
gyr_lf_z = 0
acc_x = 0
acc_y = 0
acc_z = 0
acc_z2 = 0
speed = 0
tms_cam = 0  # 时间戳 33
tms_entry = 0
tms_imu = 0
gyr_all_cnt = 0
acc_all_cnt = 0
gyr_all_node = None
acc_all_node = None
distance = 0.0

tms_time = 10000

frame_count = 0
img_data = None

clientFPS = 0


def cvtColor_YUV2NV12(_src, _dst_w=None, _dst_h=None):
    if not (_dst_w == None or _dst_h == None):
        img_new = cv2.resize(_src, (_dst_w, _dst_h))
    else:
        img_new = _src
        _dst_h = img_new.shape[0]
        _dst_w = img_new.shape[1]
    img_NV12 = np.zeros((int(_dst_h / 2 * 3), _dst_w), np.uint8)
    img_uv = np.reshape(img_new[::2, ::2, 1:], (int(_dst_h / 2), _dst_w))
    img_NV12[:_dst_h] = img_new[:, :, 0]  # Y
    img_NV12[_dst_h:] = img_uv  # UV
    return img_NV12


def Version_3_Write(frame_count, img_data):
    global time_data, tms_time, gyr_raw_x, gyr_raw_y, gyr_raw_z, gyr_lf_x, gyr_lf_y, gyr_lf_z, acc_x, acc_y, acc_z, acc_z2, speed
    tms_time += 33
    if time_data is None:  # 开启录制时，获取录制起始时间
        time_tuple = time.localtime(time.time())

        time_data = "{}{}{}_{}{}{}".format(time_tuple[0], time_tuple[1], time_tuple[2], time_tuple[3],
                                           time_tuple[4],
                                           time_tuple[5])
    else:
        if not os.path.exists('./nv12/{}'.format(time_data)):
            os.makedirs('./nv12/{}'.format(time_data))

        # 新建录制文件 文件名：int(帧号/900)
        if os.path.exists(
                './nv12/{}/{}.rawdata'.format(time_data, (int(frame_count / 901) + 1))) is False and os.path.exists(
            './nv12/{}/{}.imgdata'.format(time_data, (int(frame_count / 901) + 1))) is False:

            with open(
                    r"./nv12/{}/{}.imgdata".format(time_data, (int(frame_count / 901) + 1)), "ab+") as img_file:
                ## 输出图片
                img_file.write(struct.pack(
                    "@I", int(hex(frame_count), 16)))
                img_file.write(img_data)
            img_file.close()
            with open(
                    r"./nv12/{}/{}.rawdata".format(time_data, (int(frame_count / 901) + 1)), "a") as raw_file:
                raw_file.write('date:' + time_data + '\n')
                raw_file.write('cam_x:320,cam_y:215,cam_h:1.500000,ppd:35.6,fx:919.000000,fy:920.000000\n')
                raw_file.write('width:1280,height:720,frame_rate:30,format:NV12,img_size:1382400\n')
                raw_file.write('Odr_GYR:400,Range_GYR:125,Odr_ACC:250,Range_ACC:2\n\n')

                # 输出数据
                raw_file.write('image_frame_idx:%d,image_exists:1\n' % frame_count)
                raw_file.write(
                    'imu_valid:1,gyr.raw.x:%f,gyr.raw.y:%f,gyr.raw.z:%f,' % (gyr_raw_x, gyr_raw_y, gyr_raw_z))
                raw_file.write('gyr.lf.x:%f,gyr.lf.y:%f,gyr.lf.z:%f,' % (gyr_lf_x, gyr_lf_y, gyr_lf_z))
                raw_file.write('acc.x:%f,acc.y:%f,acc.z:%f,acc.z2:%f,' % (acc_x, acc_y, acc_z, acc_z2))
                raw_file.write('speed:%d,' % speed)
                raw_file.write('tms_cam:%d,' % tms_time)
                raw_file.write('tms_entry:%d,' % tms_time)
                raw_file.write('tms_imu:%d,' % (tms_time - 33))
                raw_file.write('gyr_all_cnt:13,acc_all_cnt:9\n')
                raw_file.write(
                    'gyr_all_node:-48,-123,-25,118,-176,-24,148,-179,-24,10,-150,-10,-168,-127,-14,-195,-124,-44,-17,-112,-58,211,-87,-32,255,-45,18,93,-23,48,-116,-37,44,-135,-90,24,32,-143,2,\n')
                raw_file.write(
                    'acc_all_node:-1025,13,-131,-1031,11,-129,-1031,13,-131,-1033,11,-131,-1035,15,-123,-1025,7,-129,-1035,15,-129,-1033,11,-131,-1037,7,-133,\n\n')

                raw_file.close()
            print("creat floder\n")
        else:
            with open(
                    r"./nv12/{}/{}.imgdata".format(time_data, (int(frame_count / 901) + 1)), "ab+") as img_file:
                ## 输出图片
                img_file.write(struct.pack(
                    "@I", int(hex(frame_count), 16)))
                img_file.write(img_data)
            with open(
                    r"./nv12/{}/{}.rawdata".format(time_data, (int(frame_count / 901) + 1)), "a") as raw_file:
                # 输出数据
                raw_file.write('image_frame_idx:%d,image_exists:1\n' % frame_count)
                raw_file.write(
                    'imu_valid:1,gyr.raw.x:%f,gyr.raw.y:%f,gyr.raw.z:%f,' % (gyr_raw_x, gyr_raw_y, gyr_raw_z))
                raw_file.write('gyr.lf.x:%f,gyr.lf.y:%f,gyr.lf.z:%f,' % (
                    gyr_lf_x, gyr_lf_y, gyr_lf_z))
                raw_file.write('acc.x:%f,acc.y:%f,acc.z:%f,acc.z2:%f,' % (
                    acc_x, acc_y, acc_z, acc_z2))
                raw_file.write('speed:%d,' % speed)
                raw_file.write('tms_cam:%d,' % tms_time)
                raw_file.write('tms_entry:%d,' % tms_time)
                raw_file.write('tms_imu:%d,' % (tms_time - 33))

                raw_file.write('gyr_all_cnt:13,acc_all_cnt:9\n')
                raw_file.write(
                    'gyr_all_node:-48,-123,-25,118,-176,-24,148,-179,-24,10,-150,-10,-168,-127,-14,-195,-124,-44,-17,-112,-58,211,-87,-32,255,-45,18,93,-23,48,-116,-37,44,-135,-90,24,32,-143,2,\n')
                raw_file.write(
                    'acc_all_node:-1025,13,-131,-1031,11,-129,-1031,13,-131,-1033,11,-131,-1035,15,-123,-1025,7,-129,-1035,15,-129,-1033,11,-131,-1037,7,-133,\n\n')

                print('Raw Frame Count: %d\n' % frame_count)


fileName = 0
all_acc_x = 0.0
all_acc_y = 0.0
all_acc_z = 0.0


def Version_4_Write():
    global time_data, tms_time, gyr_raw_x, gyr_raw_y, gyr_raw_z, gyr_lf_x, gyr_lf_y, gyr_lf_z, acc_x, acc_y, acc_z, acc_z2, speed

    global frame_count, img_data, fileName, all_acc_x, all_acc_y, all_acc_z

    tms_time += 33
    if frame_count == 0:  # 开启录制时，获取录制起始时间
        time_tuple = time.localtime(time.time())

        time_data = "{}{:0>2}{:0>2}_{:0>2}{:0>2}{:0>2}".format(time_tuple[0], time_tuple[1], time_tuple[2],
                                                               time_tuple[3],
                                                               time_tuple[4], time_tuple[5])
    else:
        if not os.path.exists('./new_QAP/{}'.format(time_data)):
            os.makedirs('./new_QAP/{}'.format(time_data))

        # 新建录制文件 文件名：int(帧号/900)  #
        if os.path.exists('./new_QAP/{}/{}.rawdata'.format(time_data, fileName)) is False and os.path.exists(
                './new_QAP/{}/{}.imgdata'.format(time_data, fileName)) is False:

            with open(
                    r"./new_QAP/{}/{}.imgdata".format(time_data, fileName),
                    "ab+") as img_file:
                ## 输出图片
                img_file.write(struct.pack(
                    "@I", int(hex(frame_count), 16)))
                img_file.write(img_data)
                print('Img Frame Count: %d' % frame_count)
            with open(
                    r"./new_QAP/{}/{}.rawdata".format(time_data, fileName),
                    "a") as raw_file:
                # 文件头数据
                raw_file.write('Date:' + time_data + ';')
                raw_file.write('Camera_Deploy:0;Camera_Position:0\n')
                raw_file.write(
                    'Width:1280;Height:720;Frame_Rate:30;Format:NV12;Image_Size:1382400;Platform:Carla\n')
                raw_file.write(
                    'Camera_x:640;Camera_y:360;Camera_h:%f;PPD:35.6;Fx:1838.0;Fy:1840.0\n\n' % cam_h)

                # 输出数据
                raw_file.write('frame_index:%d;image_exists:1;' % frame_count)
                raw_file.write('tus_cam:%d;tus_entry:%d;' % (tms_time, tms_time))
                raw_file.write('speed:%d;' % speed)
                raw_file.write('imu_exists:1;tus_imu:%d;' % (tms_time - 33))
                all_acc_x += acc_x
                all_acc_y += acc_y
                all_acc_z += acc_z

                raw_file.write('gyro:%.3f,%.3f,%.3f;' % (gyr_lf_x, gyr_lf_y, gyr_lf_z))
                raw_file.write(
                    'acc:%.1f,%.1f,%.1f\n' % (
                        all_acc_x / frame_count, all_acc_y / frame_count, all_acc_z / frame_count))
                raw_file.write('distance:%f\n\n' % distance)
                raw_file.close()

            print("creat floder\n")
        else:
            with open(
                    r"./new_QAP/{}/{}.imgdata".format(time_data, fileName),
                    "ab+") as img_file:
                ## 输出图片
                img_file.write(struct.pack(
                    "@I", int(hex(frame_count), 16)))
                img_file.write(img_data)
                # print('Img Frame Count: %d' % frame_count)
            with open(
                    r"./new_QAP/{}/{}.rawdata".format(time_data, fileName),
                    "a") as raw_file:
                # 输出数据
                raw_file.write('frame_index:%d;image_exists:1;' % frame_count)
                raw_file.write('tus_cam:%d;tus_entry:%d;' % (tms_time, tms_time))
                raw_file.write('speed:%d;' % speed)
                raw_file.write('imu_exists:1;tus_imu:%d;' % (tms_time - 33))
                all_acc_x += acc_x
                all_acc_y += acc_y
                all_acc_z += acc_z
                raw_file.write('gyro:%.3f,%.3f,%.3f;' % (gyr_lf_x, gyr_lf_y, gyr_lf_z))

                # raw_file.write(
                #     'acc:%.1f,%.1f,%.1f\n' % (
                #         all_acc_x / frame_count, all_acc_y / frame_count, all_acc_z / frame_count))

                raw_file.write(
                    'acc:%.1f,%.1f,%.1f\n' % (acc_x, acc_y, acc_z))
                raw_file.write('distance:%f\n\n' % distance)
                # print('acc:%.1f,%.1f,%.1f\n' % (
                #     all_acc_x / frame_count, all_acc_y / frame_count, all_acc_z / frame_count))
                # raw_file.close()
    fileName = int(frame_count / 900) + 1
    # print('WriteData: %d\n\n' % frame_count)
